# -*- coding: utf-8 -*-

import argparse
import os
import cv2
import math
import random
import numpy as np
from matplotlib import pyplot as plt
import utils
import preprocess
from predict import Predict

# 设置CMD参数
parser = argparse.ArgumentParser()
parser.add_argument('--image_path', type=str, help='image path')
parser.add_argument('--debug_flag', type=str, help='save debug images')
args = parser.parse_args()

## recognize
## rows
## id_cell
## score_cell
## score_number

DEBUG_FLAGS = [args.debug_flag]

def recognize(image_path):
    """
    执行成绩识别过程
    1. 图片二值化；
    2. 识别表格。分别实现图片的横线和竖线，相交后取交点；
    :param image_path:
    :return:
    """
    if not os.path.exists(image_path):
        print('File not exist. %s' % image_path)
        return None

    img = utils.open_rgb_image(image_path)

    # 提取并矫正图片中的表格区域
    img_table = preprocess.correct_table(img)
    if 'recognize' in DEBUG_FLAGS:
        utils.save_image('output/table_img.jpg', img_table)

    # 表格二值化
    img_table_thresh = preprocess.threshold(img_table, block_size=25, c=30)
    if 'recognize' in DEBUG_FLAGS:
        utils.save_image('output/img_table_thresh.jpg', img_table_thresh)    

    # 获取表格的焦点坐标
    xs, ys = get_intersections(img_table_thresh, 60)
    x_lines = filter_lines(xs)
    # y_lines = filter_lines(ys)
    # 输出校正后的表格
    if 'recognize' in DEBUG_FLAGS:
        img_table_lined = draw_lines(img_table, [], x_lines)
        utils.save_image('output/img_table_lined.jpg', img_table_lined)

    # 将表格的单元格从图片中切出
    idx = 0
    row_id = 0
    score_images = {}
    for img_row, img_i in cut_image_rows(img_table, x_lines, 0):

        # 去除表头和表格其他部分
        if img_row.shape[0] > 150:
            continue

        # 校正行图片        
        img_row_rotated, angle = preprocess.angle_correct(img_row)

        # 从行图片中识别竖线
        img_row_thresh = preprocess.threshold(img_row_rotated)

        # 为每行补充上下两条横线，用于后续交叉点的识别
        h, w = img_row_thresh.shape
        ## 上下两端绘制横线
        cv2.line(img_row_thresh, (0, 2), (w, 2), (255), 5)
        cv2.line(img_row_thresh, (0, h-6), (w, h-6), (255), 5)
        ## 左右两端绘制竖线
        cv2.line(img_row_thresh, (2, 0), (2, h), (255), 5)
        cv2.line(img_row_thresh, (w-2, 0), (w, h-6), (255), 5)

        # 获取交点坐标
        xs, ys = get_intersections(img_row_thresh, 1, 1)
        x_lines = filter_lines(xs)
        y_lines = filter_lines(ys) 
        # 输出校正后的表格
        if 'rows' in DEBUG_FLAGS:
            img_row_celled = draw_lines(img_row_rotated, y_lines, x_lines)
            utils.save_image('output/img_row_%d_inter.jpg' % (img_i), img_row_celled)

        # 切割列
        ## 提取第一列ID列
        img_cell_id = img_row_rotated[:, y_lines[0]:y_lines[1]]
        if 'rows' in DEBUG_FLAGS:
            utils.save_image('output/img_row_%d_id.jpg' % (img_i), img_cell_id)

        ## 提取最后一列成绩列
        img_cell_score = img_row_rotated[:, y_lines[-2]:y_lines[-1]]
        if 'rows' in DEBUG_FLAGS:
            utils.save_image('output/img_row_%d_score.jpg' % (img_i), img_cell_score)

        # 识别ID列和成绩列数字
        id = parse_id_number(img_cell_id, img_i)
        if id == -1.0:
            utils.save_image('output/img_row_%d_id_failed.jpg' % (img_i), img_cell_id)
            id = row_id + 1

        images = parse_score(img_cell_score, img_i)
        score_images[img_i] = (id, images)
        row_id = id

    # 计算所有数字图片宽度的中位数
    # 用于识别图片中包含一位数还是两位数
    number_image_width_list = []
    for idx, (id, images) in score_images.items():
        for i, num_img in enumerate(images): 
            number_image_width_list.append(num_img.shape[1])
    median_value = np.median(number_image_width_list) * 2
    number_image_width_list.sort()
    # plt.bar(range(len(number_image_width_list)), number_image_width_list)
    # plt.show()
    # print(median_value)

    pred = Predict()
    def predict_number(num_img):
        """
        识别手写成绩读数
        :param num_img: 成绩数字图片
        """
        # 为每个数字增加padding，和MNIST的图片格式匹配
        size = max(num_img.shape[1], num_img.shape[0]) + 40
        dw = (size - num_img.shape[1]) // 2
        dh = (size - num_img.shape[0]) // 2
        num_img_padded = cv2.copyMakeBorder(num_img, dh, dh, dw, dw, cv2.BORDER_CONSTANT, value=0)
        # utils.save_image('output/img_row_%d.jpg'%(random.randint(1, 1000)), num_img_padded)
        # 识别读数
        num = pred.predict(num_img_padded)
        return num

    for idx, (id, images) in score_images.items():
        score = 0.0
        for i, num_img in enumerate(images): 
            if num_img.shape[1] <= median_value:
                # 单个数字
                ## A*1/10^0 + B*1/10^1 + C*1/10^2
                coeff = 1 / math.pow(10, i) * 10
                num = predict_number(num_img)
                score += num * coeff
                if 'score_number' in DEBUG_FLAGS:
                    utils.save_image('output/img_row_%d_%d.jpg'%(idx ,i), num_img)
            else:
                # 两个数字
                w = num_img.shape[1] // 2
                num1_img = num_img[0:num_img.shape[0],0:w]
                num1 = predict_number(num1_img)

                num2_img = num_img[0:num_img.shape[0],w:num_img.shape[1]]
                num2 = predict_number(num2_img)

                score = num1 * 10 + num2

                if 'score_number' in DEBUG_FLAGS:
                    utils.save_image('output/img_row_%d_%d.jpg'%(idx, 0), num_img)
                    utils.save_image('output/img_row_%d_%d.jpg'%(idx, 1), num1_img)
                    utils.save_image('output/img_row_%d_%d.jpg'%(idx, 2), num2_img)

                if (len(images) == 2):
                    # 如果还有一张图片，那么应该是小数点后面的数字
                    num3 = predict_number(images[1])
                    score += 0.1 * num3
                    if 'score_number' in DEBUG_FLAGS:
                        utils.save_image('output/img_row_%d_%d.jpg'%(idx, 3), images[1])
                break
        if (score < 60.0):
            score = 0.0
        print('id=%d, score=%f' % (id, score))

    print('Recognize finished.')


def parse_score(img_cell_score, img_row_idx):

    # 二值化
    img_cell_score_thresh = preprocess.threshold(img_cell_score)
    if 'score_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_score_athresh.jpg' % (img_row_idx), img_cell_score_thresh)    

    # 通过霍夫变换识别横线，然后在横线上绘制白色线条去除横线
    deleted_lines = cv2.HoughLinesP(img_cell_score_thresh, 1, np.pi / 180, 80, img_cell_score_thresh.shape[1]*0.75, 10)
    deleted_lines = np.reshape(deleted_lines, (-1, 4)).tolist()
    for x1, y1, x2, y2 in deleted_lines:
        # 由于手写数字很容易将其中笔画识别成直线
        # 排除较倾斜的直线
        angle = abs(utils.calc_slope_angle((x1, y1, x2, y2)))
        if angle > 5 and angle + 5 < 90:
            continue
        cv2.line(img_cell_score_thresh, (x1, y1), (x2, y2), (0, 0, 0), 5)
    if 'score_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_score_clean.jpg' % (img_row_idx), img_cell_score_thresh)    

    img_cell_score_thresh_copied = img_cell_score_thresh.copy()

    # 使用形态学运算滤除噪点
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    img_cell_score_thresh = cv2.morphologyEx(img_cell_score_thresh, cv2.MORPH_OPEN, kernel, iterations=2)
    if 'score_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_score_cleaned.jpg' % (img_row_idx), img_cell_score_thresh)    

    ## 将多位数的数字横向合并起来
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 2))
    img_cell_id_thresh_merged = cv2.dilate(img_cell_score_thresh, kernel, iterations=7)
    if 'score_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_score_merged.jpg' % (img_row_idx), img_cell_id_thresh_merged)

    ## 提取成绩数字区域，提取面积最大的轮廓
    contours, _ = cv2.findContours(img_cell_id_thresh_merged, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    outer_contour = max(contours, key=lambda cnt: cv2.contourArea(cnt))
    x, y, w, h = cv2.boundingRect(outer_contour)
    img_cell_score_thresh_outer = img_cell_score_thresh_copied[y:y+h,x:x+w] 
    img_cell_score_thresh_outer = cv2.copyMakeBorder(img_cell_score_thresh_outer,20,20,10,10,cv2.BORDER_CONSTANT,value=0)
    if 'score_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_score_outer_a.jpg' % (img_row_idx), img_cell_score_thresh_outer)

    score_images = []
    # 根据x坐标排序所有轮廓
    img_cell_score_thresh_copied = img_cell_score_thresh_outer.copy()
    digit_contours, _ = cv2.findContours(img_cell_score_thresh_outer, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    digit_contours.sort(key=lambda cnt: cv2.boundingRect(cnt)[0])
    for ctn in digit_contours:
        x, y, w, h = cv2.boundingRect(ctn) 
        if w * h < 500 or w / h >= 5.0:
            continue
        img_cell_score_thresh_cropped = img_cell_score_thresh_outer[y:y+h,x:x+w]
        if 'score_number' in DEBUG_FLAGS:
            utils.save_image('output/img_row_%d_cropped_%d.jpg'%(img_row_idx, x), img_cell_score_thresh_cropped)
        score_images.append(img_cell_score_thresh_cropped)
        cv2.rectangle(img_cell_score_thresh_copied, (x, y), (x+w, y+h), (255, 255, 255), 2)
    if 'score_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_score_outer_bounded.jpg' % (img_row_idx), img_cell_score_thresh_copied)    
    return score_images 


def parse_id_number(img_cell_id, img_row_idx):
    """
    解析ID列序号
    :param img_cell_id: id列单元格图片
    :param img_row_idx: 图片所在行号
    :return: 返回id列序号
    """
    ## 识别ID列编号
    img_cell_id_thresh = preprocess.threshold(img_cell_id)
    if 'id_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_id_athresh.jpg' % (img_row_idx), img_cell_id_thresh)

    ## 霍夫变换，去除图片中多余的表格边框线
    deleted_lines = cv2.HoughLinesP(img_cell_id_thresh, 1, np.pi / 180, 80, 200, 10)
    deleted_lines = np.reshape(deleted_lines, (-1, 4)).tolist()
    for x1, y1, x2, y2 in deleted_lines:
        cv2.line(img_cell_id_thresh, (x1, y1), (x2, y2), (0, 0, 0), 5)
    if 'id_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_id_clean.jpg' % (img_row_idx), img_cell_id_thresh)

    # 使用形态学运算滤除噪点
    img_cell_id_thresh_tpl = img_cell_id_thresh.copy()
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
    img_cell_id_thresh_tpl = cv2.morphologyEx(img_cell_id_thresh_tpl, cv2.MORPH_OPEN, kernel, iterations=2)
    if 'id_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_id.jpg' % (img_row_idx), img_cell_id_thresh_tpl)

    ## 将多位数的数字横向合并起来
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 1))
    img_cell_id_thresh_tpl = cv2.dilate(img_cell_id_thresh_tpl, kernel, iterations=10)
    if 'id_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_id_merged.jpg' % (img_row_idx), img_cell_id_thresh_tpl)

    ## 提取序号数字
    img_cell_num_thresh = None
    contours, _ = cv2.findContours(img_cell_id_thresh_tpl, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  
    for ctn in contours:
        x, y, w, h = cv2.boundingRect(ctn) 
        if w * h < 1000 or w / h >= 5.0: 
            continue
        img_cell_num_thresh = img_cell_id_thresh[y:y+h+5,x:x+w+5]
        ## 将数字图片上下左右各增加20的Padding
        size = max(img_cell_num_thresh.shape[1], img_cell_num_thresh.shape[0]) + 40
        dw = (size - img_cell_num_thresh.shape[1]) // 2
        dh = (size - img_cell_num_thresh.shape[0]) // 2
        img_cell_num_thresh = cv2.copyMakeBorder(img_cell_num_thresh,dh,dh,dw,dw,cv2.BORDER_CONSTANT,value=0)
        break
    if img_cell_num_thresh is None:
        return -1.0
    if 'id_cell' in DEBUG_FLAGS:
        utils.save_image('output/img_row_%d_id_padding.jpg' % (img_row_idx), img_cell_num_thresh)    

    ## 将二值化图片反转为白底黑字
    idstr = utils.recognize_id(~img_cell_num_thresh)

    ## 如果为数字才返回序号
    if idstr.isdigit():
        return int(idstr)
    return -1.0


def cut_image_rows(img_rotated, x_lines, margin=10):
    """
    从原图中切出行
    :param img_rotated: 校正后的图片
    :param x_lines: x坐标
    :param margin: 切图的边距
    :return: 每行返回一次，返回OpenCV图片和行坐标
    """
    for xi in range(len(x_lines)-1):
        x0, x1 = x_lines[xi], x_lines[xi+1]
        img_row = img_rotated[x0-margin:x1+margin, :]
        if utils.__DEBUG:
            cv2.imwrite('output/%d.jpg' % xi, img_row)
        yield img_row, xi


def cut_image(img_rotated, x_lines, y_lines, margin=10):
    """
    从原图中切出单元格
    :param img_rotated: 校正后的图片
    :param x_lines: x坐标
    :param y_lines: y坐标
    :param margin: 切图的边距
    :return: 每个单元格返回一次
    """
    for xi in range(len(x_lines)-1):
        for yi in range(len(y_lines)-1):
            x0, y0 = x_lines[xi], y_lines[yi]
            x1, y1 = x_lines[xi+1], y_lines[yi+1]
            img_cell = img_rotated[x0:x1, y0:y1]
            if utils.__DEBUG:
                cv2.imwrite('output/%d_%d.jpg' % (yi, xi), img_cell)
            yield img_cell, xi, yi


def get_intersections(img_thresh, row_precision=90, col_precision=45):
    """
    取图片横线和竖线的交点
    :param img_thresh: 二值化图片
    :param row_precision: 横线提取精度
    :param col_precision: 竖线提取进度
    :return: 返回焦点坐标数组，分为x数组和y数组
    """
    row_lines = preprocess.recognize_rows(img_thresh, row_precision)
    col_lines = preprocess.recognize_cols(img_thresh, col_precision)
    intersections = cv2.bitwise_and(row_lines, col_lines)
    return np.where(intersections > 0)


def filter_lines(lines, min_gap=50):
    """
    过滤距离太接近的线
    :param lines: 线条坐标
    :param min_gap: 线条之前的最小距离
    :return: 返回过滤后的线条坐标数组
    """
    line_list = []
    lines = np.sort(lines)
    for i in range(len(lines) - 1):
        if lines[i + 1] - lines[i] > min_gap:
            line_list.append(lines[i])
    line_list.append(lines[len(lines) - 1])
    return line_list


def draw_lines(img, x_lines, y_lines):
    """
    图片中绘制表格线、交点以及交点坐标
    :param img: 图片
    :param x_lines: 竖线
    :param y_lines: 横线
    :return: 返回绘制后的OpenCV图像
    """
    img_output = img.copy()
    h, w, c = img.shape
    for x in x_lines:
        cv2.line(img_output, (x, 0), (x, h), (0, 255, 255), 3)
    for y in y_lines:
        cv2.line(img_output, (0, y), (w, y), (255, 0, 255), 3)
    for xi in range(len(x_lines)):
        for yi in range(len(y_lines)):
            x, y = x_lines[xi], y_lines[yi]
            cv2.putText(img_output, '(%d,%d)' % (x, y), (x-100, y+30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255))
            cv2.circle(img_output, (x, y), 8, (0, 255, 0), 10)
    return img_output


if __name__ == '__main__':
    utils.__DEBUG = False
    recognize(args.image_path)
